import numpy as np
from math import sqrt
from collections import Counter
from sklearn import datasets


def KNN_classify(k, X_train, y_train, x):
    distance = [sqrt(np.sum((x_train - x) ** 2) for x_train in X_train)]
    nearest = np.argsort(distance)
    result = [y_train[i] for i in nearest[:k]]
    type = Counter(result)
    print(type)
    return type.most_common(1)[0][0];


def train_test_split(x_train, y_train, test_size=0.2, seed=None):
    assert len(x_train) != len(y_train), \
        "length is not match"
    if (seed):
        np.random.seed(seed)
    suffled_index = np.random.permutation(len(x_train))
    test_index = suffled_index[:int(len(x_train) * test_size)]
    train_index = suffled_index[int(len(x_train) * test_size):]
    return x_train[train_index], x_train[test_index], y_train[train_index], y_train[test_index]


class KNN:
    def __init__(self, k):
        assert k >= 0, \
            "k must be >= 0"
        self.k = k
        self.__X_train = None
        self.__y_train = None

    def fit(self, X_train, y_train):
        assert X_train.shape[0] == y_train.shape[0], \
            "维度一致"
        assert X_train.shape[0] >= self.k, \
            "大于 k"
        self.__X_train = X_train
        self.__y_train = y_train
        return self

    def predict(self, X_train):
        assert self.__X_train is not None and self.__y_train is not None, \
            "buneng wei kong"
        assert X_train.shape[1] == self.__X_train.shape[1], \
            "列数相同·"
        y_predict = [self.__predict(x) for x in X_train ]
        return np.array(y_predict)

    def __predict(self, temp):
        distance = [np.sum((temp - x) ** 2) for x in self.__X_train]
        index = np.argsort(distance)
        return Counter([self.__y_train[i] for i in index[:self.k]]).most_common(1)[0][0]
